package com.apisix.service.util;

import org.bouncycastle.util.io.pem.PemObject;
import org.bouncycastle.util.io.pem.PemWriter;

import java.io.IOException;
import java.io.StringWriter;
import java.security.*;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

public class KeyGeneratorUtil {

    public static Map<String, String> generateRSAKeyPair() throws NoSuchAlgorithmException, IOException {
        KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
        keyGen.initialize(2048);
        KeyPair keyPair = keyGen.generateKeyPair();

        PublicKey publicKey = keyPair.getPublic();
        PrivateKey privateKey = keyPair.getPrivate();

        String publicKeyPEM = convertToPEM(publicKey.getEncoded(), "PUBLIC KEY");
        String privateKeyPEM = convertToPEM(privateKey.getEncoded(), "PRIVATE KEY");

        Map<String, String> keyPairMap = new HashMap<>();
        keyPairMap.put("publicKey", publicKeyPEM);
        keyPairMap.put("privateKey", privateKeyPEM);

        return keyPairMap;
    }

    private static String convertToPEM(byte[] keyBytes, String description) throws IOException {
        PemObject pemObject = new PemObject(description, keyBytes);
        StringWriter stringWriter = new StringWriter();
        try (PemWriter pemWriter = new PemWriter(stringWriter)) {
            pemWriter.writeObject(pemObject);
        }
        return stringWriter.toString();
    }

    public static String convertPublicKeyToBase64(String pemKey) {
        String publicKeyPEM = pemKey
                .replace("-----BEGIN PUBLIC KEY-----", "")
                .replace("-----END PUBLIC KEY-----", "")
                .replaceAll("\\s", "");
        return publicKeyPEM;
    }

    public static String convertPrivateKeyToBase64(String pemKey) {
        String privateKeyPEM = pemKey
                .replace("-----BEGIN PRIVATE KEY-----", "")
                .replace("-----END PRIVATE KEY-----", "")
                .replaceAll("\\s", "");
        return privateKeyPEM;
    }
} 